#!/usr/local/python3/bin/python3
import numpy as np
import user_vec as uv
import time

mid = np.load("films_mid.npy", allow_pickle="TRUE")
vec = np.load("films_vec.npy", allow_pickle="TRUE")

def get_topK(uid, K=6):
    distance = np.array([])
    col_mid = uv.get_sql(uid)  # 返回已经收藏的电影
    user_vec = uv.compute_vec(col_mid)

    start_time = time.time()

    for i in range(len(vec)):
        if mid[i] in col_mid:
            distance = np.append(distance, 1000)
        else:
            tmp = uv.euclidean_distance_by_tf(vec[i], user_vec)
            distance = np.append(distance, tmp)
    top6 = np.argpartition(distance, K)[:K]  # 下标
    end_time = time.time()
    return (top6, end_time-start_time)  # 序号 and cost_time

# 序号转换成电影ID
def convert_to_mid(topk):
    res = []
    for each in topk:
        key = mid[each]
        res.append(key)
    return tuple(res)

if __name__ == '__main__':
    import sys
    # uid = int(sys.argv[1])
    uid = 11

    top6, cost_time = get_topK(uid)
    films = np.load("films_list.npy", allow_pickle="TRUE")
    print(convert_to_mid(top6))
    print(cost_time)
    for each in top6:
        film = films[each]
        print(film)
